#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
eBPF 捕获服务，用于记录网络流量
支持 BCC 和 CO-RE (Compile Once – Run Everywhere) 两种方式
"""

import logging
import os
import time
import threading
import json
from datetime import datetime
from bson import ObjectId
from db.mongo import get_db
from intercept.models.intercept import InterceptRequestRecordModel
from intercept.services.intercept_service import InterceptService

logger = logging.getLogger(__name__)

# 尝试导入 libbpf 捕获服务
try:
    from intercept.services.libbpf_capture import LibbpfCaptureService
    LIBBPF_AVAILABLE = True
except ImportError:
    LIBBPF_AVAILABLE = False
    logger.warning("libbpf 捕获服务不可用，将使用 BCC 或模拟模式")

class EbpfCaptureService:
    """eBPF 捕获服务"""

    # 捕获线程
    _capture_thread = None

    # 是否正在运行
    _is_running = False

    # 要捕获的服务列表
    _services = []

    @staticmethod
    def start_capture(interface='eth0', services=None):
        """
        启动捕获

        Args:
            interface: 网络接口
            services: 要捕获的服务列表，格式为 [{"service_id": "xxx", "service_name": "xxx", "service_ip": "xxx", "service_port": xxx}]

        Returns:
            bool: 是否成功启动
        """
        if EbpfCaptureService._is_running:
            logger.warning("捕获服务已经在运行")
            return False

        # 检查 BCC 库是否可用
        if not EbpfCaptureService.is_bcc_available():
            logger.error("BCC库不可用，无法启动捕获服务")
            return False

        # 检查服务列表
        if not services or len(services) == 0:
            logger.warning("没有指定要捕获的服务，将不会捕获任何流量")
            # 返回 True，因为这不是一个错误状态

        # 保存服务列表
        EbpfCaptureService._services = services or []

        # 尝试使用 libbpf (CO-RE) 方式
        if LIBBPF_AVAILABLE and LibbpfCaptureService.is_available():
            logger.info("使用 libbpf (CO-RE) 方式启动捕获服务")
            return LibbpfCaptureService.start_capture(interface, services)

        # 启动捕获线程
        EbpfCaptureService._is_running = True
        EbpfCaptureService._capture_thread = threading.Thread(
            target=EbpfCaptureService._capture_worker,
            args=(interface, EbpfCaptureService._services),
            daemon=True
        )
        EbpfCaptureService._capture_thread.start()

        logger.info(f"已启动捕获服务，网络接口: {interface}，服务数量: {len(EbpfCaptureService._services)}")
        return True

    @staticmethod
    def stop_capture():
        """
        停止捕获

        Returns:
            bool: 是否成功停止
        """
        if not EbpfCaptureService._is_running:
            logger.warning("捕获服务未在运行")
            return False

        # 如果使用的是 libbpf 方式，则调用 libbpf 的停止方法
        if LIBBPF_AVAILABLE and LibbpfCaptureService.is_running():
            logger.info("停止 libbpf 捕获服务")
            return LibbpfCaptureService.stop_capture()

        # 停止 BCC 捕获线程
        EbpfCaptureService._is_running = False
        if EbpfCaptureService._capture_thread:
            EbpfCaptureService._capture_thread.join(timeout=5)
            EbpfCaptureService._capture_thread = None

        logger.info("已停止捕获服务")
        return True

    @staticmethod
    def is_running():
        """
        检查是否正在运行

        Returns:
            bool: 是否正在运行
        """
        # 如果 libbpf 捕获服务正在运行，也返回 True
        if LIBBPF_AVAILABLE and LibbpfCaptureService.is_running():
            return True
        return EbpfCaptureService._is_running

    @staticmethod
    def get_services():
        """
        获取正在捕获的服务列表

        Returns:
            list: 服务列表
        """
        # 如果 libbpf 捕获服务正在运行，使用它的服务列表
        if LIBBPF_AVAILABLE and LibbpfCaptureService.is_running():
            return LibbpfCaptureService.get_services()
        return EbpfCaptureService._services

    # 检查 BCC 库是否可用
    @staticmethod
    def is_bcc_available():
        """检查 BCC 库是否可用"""
        try:
            from bcc import BPF
            logger.info("BCC库加载成功，可以使用 eBPF 功能")
            return True
        except ImportError:
            logger.warning("BCC库加载失败，无法使用 eBPF 功能")
            return False

    @staticmethod
    def _capture_worker(interface, services):
        """
        捕获工作线程

        Args:
            interface: 网络接口
            services: 要捕获的服务列表
        """
        # 检查 BCC 库是否可用
        if not EbpfCaptureService.is_bcc_available():
            logger.error("BCC库不可用，无法启动捕获服务")
            return False

        logger.info(f"开始使用 eBPF 捕获网络接口 {interface} 上的流量")
        service_info = []
        for s in services:
            name = s.get('service_name', 'unknown')
            ip = s.get('service_ip', '0.0.0.0')
            port = s.get('service_port', 0)
            service_info.append(f"{name}({ip}:{port})")
        logger.info(f"监控服务列表: {service_info}")

        try:
            # 导入 BCC 库
            from bcc import BPF

            # 加载 eBPF 程序
            bpf_text = """
            #include <uapi/linux/ptrace.h>
            #include <net/sock.h>
            #include <bcc/proto.h>
            #include <linux/ip.h>
            #include <linux/tcp.h>

            // 定义数据结构
            struct http_data_t {
                u64 ts;                // 时间戳
                u32 pid;               // 进程ID
                u32 saddr;             // 源IP
                u32 daddr;             // 目标IP
                u16 sport;             // 源端口
                u16 dport;             // 目标端口
                u8 direction;          // 方向：1=请求，2=响应
                char data[4096];       // 数据
            };

            // 定义事件表
            BPF_PERF_OUTPUT(http_events);

            // 处理 TCP 数据
            int trace_tcp_data(struct pt_regs *ctx, struct sock *sk, struct sk_buff *skb) {
                if (skb == NULL || sk == NULL)
                    return 0;

                // 获取 TCP 头
                struct iphdr *iph = (struct iphdr *)skb->head;
                struct tcphdr *th = (struct tcphdr *)(skb->head + sizeof(struct iphdr));

                // 检查是否是 HTTP 流量
                u16 dport = 0;
                u16 sport = 0;
                bpf_probe_read(&dport, sizeof(dport), &th->dest);
                bpf_probe_read(&sport, sizeof(sport), &th->source);
                dport = ntohs(dport);
                sport = ntohs(sport);

                // 只处理 HTTP 流量
                if (dport != 80 && sport != 80 && dport != 443 && sport != 443)
                    return 0;

                // 获取数据
                char *data = (char *)(skb->head + sizeof(struct iphdr) + sizeof(struct tcphdr));

                // 创建事件
                struct http_data_t event = {};
                event.ts = bpf_ktime_get_ns();
                event.pid = bpf_get_current_pid_tgid() >> 32;
                event.saddr = iph->saddr;
                event.daddr = iph->daddr;
                event.sport = sport;
                event.dport = dport;
                event.direction = (dport == 80 || dport == 443) ? 1 : 2;
                bpf_probe_read(event.data, sizeof(event.data), data);

                // 提交事件
                http_events.perf_submit(ctx, &event, sizeof(event));

                return 0;
            }
            """

            # 编译 eBPF 程序
            logger.info("开始编译 eBPF 程序...")
            try:
                b = BPF(text=bpf_text)
                logger.info("eBPF 程序编译成功")
            except Exception as e:
                if "Failed to compile BPF module" in str(e) and "Unable to find kernel headers" in str(e):
                    logger.error("编译 eBPF 程序失败: 无法找到内核头文件，将使用模拟模式运行")
                    # 模拟模式下的工作逻辑
                    while EbpfCaptureService._is_running:
                        time.sleep(5)  # 每5秒检查一次
                    return
                else:
                    logger.error(f"编译 eBPF 程序失败: {e}")
                    raise

            # 附加到 TCP 发送和接收函数
            logger.info("开始附加 kprobe 到 TCP 函数...")
            try:
                b.attach_kprobe(event="tcp_sendmsg", fn_name="trace_tcp_data")
                b.attach_kprobe(event="tcp_recvmsg", fn_name="trace_tcp_data")
                logger.info("kprobe 附加成功，eBPF 程序已完全启动")
            except Exception as e:
                logger.error(f"kprobe 附加失败: {e}")
                # 模拟模式下的工作逻辑
                EbpfCaptureService._run_mock_mode(services)
                return

            # 定义事件处理函数
            def process_event(cpu, data, size):
                event = b["http_events"].event(data)

                # 转换 IP 地址
                import socket
                saddr = socket.inet_ntoa(bytes.fromhex(format(event.saddr, '08x')))
                daddr = socket.inet_ntoa(bytes.fromhex(format(event.daddr, '08x')))

                # 检查是否匹配服务
                if not EbpfCaptureService._filter_service(saddr, daddr, event.sport, event.dport, services):
                    return

                # 解析 HTTP 数据
                data = event.data.decode('utf-8', errors='ignore')

                # 处理 HTTP 请求或响应
                if event.direction == 1:  # 请求
                    EbpfCaptureService._process_http_request(saddr, daddr, event.sport, event.dport, data, services)
                else:  # 响应
                    EbpfCaptureService._process_http_response(saddr, daddr, event.sport, event.dport, data, services)

            # 打开事件管道
            b["http_events"].open_perf_buffer(process_event)

            # 循环处理事件
            while EbpfCaptureService._is_running:
                try:
                    b.perf_buffer_poll(timeout=100)
                except KeyboardInterrupt:
                    break

            logger.info("捕获线程已退出")

        except ImportError:
            logger.error("未安装 BCC 库，无法使用 eBPF 功能")
            EbpfCaptureService._is_running = False

        except Exception as e:
            logger.error(f"捕获线程异常: {e}")
            EbpfCaptureService._is_running = False

    @staticmethod
    def _filter_service(saddr, daddr, sport, dport, services):
        """
        过滤服务

        Args:
            saddr: 源IP
            daddr: 目标IP
            sport: 源端口
            dport: 目标端口
            services: 服务列表

        Returns:
            bool: 是否匹配
        """
        if not services:
            return True

        for service in services:
            service_ip = service.get("service_ip")
            service_port = service.get("service_port")

            if not service_ip or not service_port:
                continue

            # 检查是否匹配服务
            if (saddr == service_ip and sport == service_port) or (daddr == service_ip and dport == service_port):
                return True

        return False

    @staticmethod
    def _process_http_request(saddr, daddr, sport, dport, data, services):
        """
        处理 HTTP 请求

        Args:
            saddr: 源IP
            daddr: 目标IP
            sport: 源端口
            dport: 目标端口
            data: HTTP 数据
            services: 服务列表
        """
        try:
            # 解析 HTTP 请求
            lines = data.split('\r\n')
            if not lines:
                return

            # 解析请求行
            request_line = lines[0].split(' ')
            if len(request_line) < 3:
                return

            method = request_line[0]
            path = request_line[1]
            version = request_line[2]

            # 解析请求头
            headers = {}
            for i in range(1, len(lines)):
                line = lines[i]
                if not line:
                    break

                parts = line.split(':', 1)
                if len(parts) == 2:
                    headers[parts[0].strip()] = parts[1].strip()

            # 解析请求体
            body = None
            if 'Content-Length' in headers:
                content_length = int(headers['Content-Length'])
                if content_length > 0:
                    body_start = data.find('\r\n\r\n') + 4
                    if body_start > 0 and body_start < len(data):
                        body = data[body_start:body_start + content_length]

            # 构建请求 URL
            host = headers.get('Host', f"{daddr}:{dport}")
            url = f"http://{host}{path}"

            # 查找匹配的服务
            service_id = None
            service_name = None
            for service in services:
                if service.get("service_ip") == daddr and service.get("service_port") == dport:
                    service_id = service.get("service_id")
                    service_name = service.get("service_name")
                    break

            if not service_id:
                return

            # 保存请求信息
            EbpfCaptureService._save_request(service_id, service_name, method, url, headers, body)

        except Exception as e:
            logger.error(f"处理 HTTP 请求异常: {e}")

    @staticmethod
    def _process_http_response(saddr, daddr, sport, dport, data, services):
        """
        处理 HTTP 响应

        Args:
            saddr: 源IP
            daddr: 目标IP
            sport: 源端口
            dport: 目标端口
            data: HTTP 数据
            services: 服务列表
        """
        try:
            # 解析 HTTP 响应
            lines = data.split('\r\n')
            if not lines:
                return

            # 解析状态行
            status_line = lines[0].split(' ', 2)
            if len(status_line) < 3:
                return

            version = status_line[0]
            status = int(status_line[1])
            reason = status_line[2]

            # 解析响应头
            headers = {}
            for i in range(1, len(lines)):
                line = lines[i]
                if not line:
                    break

                parts = line.split(':', 1)
                if len(parts) == 2:
                    headers[parts[0].strip()] = parts[1].strip()

            # 解析响应体
            body = None
            if 'Content-Length' in headers:
                content_length = int(headers['Content-Length'])
                if content_length > 0:
                    body_start = data.find('\r\n\r\n') + 4
                    if body_start > 0 and body_start < len(data):
                        body = data[body_start:body_start + content_length]

            # 查找匹配的服务
            service_id = None
            for service in services:
                if service.get("service_ip") == saddr and service.get("service_port") == sport:
                    service_id = service.get("service_id")
                    break

            if not service_id:
                return

            # 保存响应信息
            EbpfCaptureService._save_response(service_id, status, headers, body)

        except Exception as e:
            logger.error(f"处理 HTTP 响应异常: {e}")

    @staticmethod
    def _save_request(service_id, service_name, method, url, headers, body):
        """
        保存请求信息

        Args:
            service_id: 服务ID
            service_name: 服务名称
            method: 请求方法
            url: 请求URL
            headers: 请求头
            body: 请求体
        """
        try:
            # 创建临时请求记录
            db = get_db()
            db.temp_http_requests.insert_one({
                "service_id": service_id,
                "service_name": service_name,
                "method": method,
                "url": url,
                "headers": headers,
                "body": body,
                "created_at": datetime.now()
            })

            logger.debug(f"已保存请求: {method} {url}")

        except Exception as e:
            logger.error(f"保存请求信息异常: {e}")



    @staticmethod
    def _save_response(service_id, status, headers, body):
        """
        保存响应信息

        Args:
            service_id: 服务ID
            status: 状态码
            headers: 响应头
            body: 响应体
        """
        try:
            # 查找最近的请求记录
            db = get_db()
            request = db.temp_http_requests.find_one(
                {"service_id": service_id},
                sort=[("created_at", -1)]
            )

            if not request:
                return

            # 删除临时请求记录
            db.temp_http_requests.delete_one({"_id": request["_id"]})

            # 创建完整的请求记录
            InterceptService.create_or_update_request_record(
                service_id=service_id,
                method=request["method"],
                url=request["url"],
                request_headers=request["headers"],
                request_body=request["body"],
                response_status=status,
                response_headers=headers,
                response_body=body
            )

            logger.debug(f"已保存响应: {request['method']} {request['url']} -> {status}")

        except Exception as e:
            logger.error(f"保存响应信息异常: {e}")
